import xlrd
import numpy as np
import scipy.io as sio
from matplotlib import pyplot as plt
import time
import Ground_State_Initial_Settings
import sys
sys.path.append('D:\\vc\\Rb87Python\\')
import basic_function.Rb87BEC_Basic_Functions as bf
import os
import socket

os.system('chcp 65001')
Host_Name = socket.gethostname()
if Host_Name == 'VictorChengPC':
    Computer_Name = 'X:\\Work\\mail.ustc.edu.cn\\Rb87BEC - Files\\'
elif Host_Name == 'calcenter2':
    Computer_Name = 'D:\\vc\\mail.ustc.edu.cn\\Rb87BEC - Files\\'

# Index = 0.05

# %%
time_start = time.time()

# %% Switches
FitBEC_Switch = 0
FitDraw_Switch = 0
OriImage_Switch = 1
Data_Total = []
# %% Modes
File_Mode = 1
Atom_Mode = 0

# %% Initial settings
'''
No need to -1 on row number.
'''
Mat_Data = [51]
# Mat_Data = np.linspace(3, 24, 22).astype(int)
Data_Exl = xlrd.open_workbook(Computer_Name + r'Rb87files\3DSOC_ground_state\3DSOC_ground_state.xlsx')
Table = Data_Exl.sheets()[0]
# Matrix_Shift = [0, 0]
clims = [-0.1, 0.5]
# Wait_Sec = 1
# Photo_Name_Tag='22860test_'
# Photo_Total = ''
# if File_Mode == 0:
#    PhotoName_Pre = 'ascOD'+Photo_Name_Tag
#    PhotoName_End = '.asc'
#    
# elif File_Mode == 1:
#    PhotoName_Pre = Photo_Name_Tag
#    PhotoName_End = 'od.mat'

# %% Calculate all photos
Data_Every_Row = {}
AtomClouds_Every_Row = {}
for Row_Num in range(len(Mat_Data)):
    Data_This_Row = {}
    Exl_RowNum = Mat_Data[Row_Num] - 1
    Input_Matrix = Table.row_values(Exl_RowNum)
    Parameters = {}

    # Set parameters
    YearMonthDay = str(int(Input_Matrix[0]))
    Parameters['YearMonthDay'] = YearMonthDay
    Ini_Photo = int(Input_Matrix[1])
    Parameters['Ini_Photo'] = Ini_Photo
    Fin_Photo = int(Input_Matrix[2])
    Parameters['Fin_Photo'] = Fin_Photo
    Repeat = int(Input_Matrix[3])
    Parameters['Repeat'] = Repeat
    V0 = Input_Matrix[4]
    Parameters['V0'] = V0
    Vy = Input_Matrix[5]
    Parameters['Vy'] = Vy
    Detuning = Input_Matrix[6]
    Parameters['Detuning'] = Detuning
    Cam = Input_Matrix[7]
    M0 = Input_Matrix[8]
    Parameters['M0'] = M0

    # # @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
    # Skip_Start = Input_Matrix[10]
    # Skip_End = Input_Matrix[11]
    # # @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@

    #    Enclose_Switch = Input_Matrix[9]
    #    Enclose_Switch = ['Circle', 'Circle']
    Photo_Name_Tag = Cam + 'test_'
    Photo_Total = ''
    if File_Mode == 0:
        PhotoName_Pre = 'ascOD' + Photo_Name_Tag
        PhotoName_End = '.asc'
    elif File_Mode == 1:
        PhotoName_Pre = Photo_Name_Tag
        PhotoName_End = 'od.mat'
    Parameters_str = str(Parameters)
    Parameters_str = Parameters_str.replace('{', '')
    Parameters_str = Parameters_str.replace('}', '')
    Parameters_str = Parameters_str.replace(' ', '')
    Parameters_str = Parameters_str.replace(':', '=')
    Photo = str(Ini_Photo) + 'To' + str(Fin_Photo)
    Photo_Total += Photo + ';'
    # Data matrices
    # if Row_Num == 0:
    #     Time = np.linspace(t_Ini, t_End, (t_End-t_Ini)/Time_Step+1)

    # Addresses
    if File_Mode == 0:
        Folder_Name = r'\\CAMERA-PC\ShareFolder' + '\\' + \
                      YearMonthDay + 'morning\\' + Cam + '\\'
    elif File_Mode == 1:
        if Cam == '22834':
            #            Folder_Name = Computer_Name+'oriData'+'\\'+YearMonthDay+'New_Method\\'
            Folder_Name = '\\\\rb87bec3\\' + 'oriData' + '\\' + YearMonthDay + 'New_Method\\'
        else:
            Folder_Name = '\\\\calcenter2\\' + 'oriData' + '\\' + YearMonthDay + '\\'
    SaveFigure_Address = '\\\\calcenter2\\' + r'Rb87figures\3DSOC_ground_state\figures' + '\\' + \
        YearMonthDay + '\\' + Photo + '\\'
    os.system('mkdir ' + SaveFigure_Address)
    SaveData_Address = Computer_Name + 'Rb87Data\\3DSOC_ground_state' + '\\' + \
        YearMonthDay + '\\' + Photo + '\\'
    os.system('mkdir ' + '\"' + SaveData_Address + '\"')
    # Save_Name = YearMonthDay + Parameters_str + '3DSOC_ground_state' + Photo
    Atom_Center = Ground_State_Initial_Settings.AtomCenter_Pos(Cam)
    Cut_Width = Ground_State_Initial_Settings.CutWidths(Cam)
    Enclose_Switch = ['Circle' for ii in range(len(Cut_Width) - 1)]
    Up_Regions = []
    Dn_Regions = []
    Atom_Clouds = []
    # Calculating
    Photo_Ind = Ini_Photo
    while Photo_Ind <= Fin_Photo:
        # # @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
        # if Skip_End >= Photo_Ind >= Skip_Start:
        #     Photo_Ind += 1
        #     continue
        # # @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
        plt.close('all')
        Load_Name = Folder_Name + PhotoName_Pre + '%03d' % Photo_Ind + PhotoName_End
        Save_Name = Folder_Name + PhotoName_Pre + '%03d' % Photo_Ind + 'od'
        if Cam == '22860':
            All_x = [0, 1024]
            All_y = [0, 1024]
            # rot 43
            Up_x = [0, 1024]
            Up_y = [0, 1024]
            # rot -2
            # Up_x = [340, 500]
            # Up_y = [395, 555]
            Cut_Aera_Up = (Up_x, Up_y)
            # rot 43
            Dn_x = [430, 590]
            Dn_y = [350, 510]
            # rot -2
            # Dn_x = [490, 650]
            # Dn_y = [375,535]
            Cut_Aera_Dn = (Dn_x, Dn_y)
            Data_ThisPic_AtomNum, Data_ThisPic_AtomRegion, ne = bf.Load_Pic_And_Cal(
                Photo_Ind, Load_Name, File_Mode, Save_Name, SaveFigure_Address + 'All_' + Cam + '_',
                Atom_Center, Cut_Width, len(Cut_Width) - 1, OriImage_Switch,
                Rot_Ang=43, Axis='off', Enclose=Enclose_Switch)
            if Data_ThisPic_AtomNum == 'Error':
                print(ne)
                break
            Data_This_Row[Photo_Ind] = (Data_ThisPic_AtomNum, Data_ThisPic_AtomRegion)
            Atom_Clouds.append(np.array(Data_ThisPic_AtomNum))
            Data_ThisPic_AtomNum, Data_ThisPic_AtomRegion, Up_Region = bf.Load_Pic_And_Cal(
                Photo_Ind, Load_Name, File_Mode, Save_Name, SaveFigure_Address + 'Up_' + Cam + '_',
                Atom_Center, Cut_Width, len(Cut_Width) - 1, OriImage_Switch,
                Rot_Ang=43, Axis='off', Cut_Aera=Cut_Aera_Up, Enclose=Enclose_Switch)
            Up_Regions.append(Up_Region)
            # plt.plot((0, 1000), (440, 440), 'k-', linewidth=5)
            # plt.plot((326, 326), (0, 1000), 'k-',  linewidth=5)
            # Data_ThisPic_AtomNum, Data_ThisPic_AtomRegion, Dn_Region = bf.Load_Pic_And_Cal(
            #     Photo_Ind, Load_Name, File_Mode, Save_Name, SaveFigure_Address+'Down_'+Cam+'_',
            #     Atom_Center, Cut_Width, len(Cut_Width)-1, OriImage_Switch,
            #     Rot_Ang=43, Axis='off', Cut_Aera=Cut_Aera_Dn, Enclose=Enclose_Switch)
            Dn_Regions.append(Up_Region)
            # plt.plot((5, 100),(82, 82) , 'k-', linewidth=5)
            # plt.plot((25, 25), (5, 100), 'k-',  linewidth=5)
            figsize = (10, 10)
            Top_Up = 1.0
            Bottom_Up = 0.047
            Left_Up = 0.047
            Right_Up = 0.981
            Hspace_Up = 0.2
            Wspace_Up = 0.2
            Top_Dn = 1.0
            Bottom_Dn = 0.047
            Left_Dn = 0.047
            Right_Dn = 0.981
            Hspace_Dn = 0.2
            Wspace_Dn = 0.2
        elif Cam == '4552':
            All_x = [0, 512]
            All_y = [0, 512]
            Up_x = [150, 330]
            Up_y = [320, 500]
            Dn_x = [240, 400]
            Dn_y = [170, 330]
        elif Cam == '22834':
            All_x = [0, 1024]
            All_y = [0, 1024]
            Up_x = [0, 1024]
            Up_y = [0, 1024]
            Cut_Aera_Up = (Up_x, Up_y)
            Dn_x = [340, 480]
            Dn_y = [395, 625]
            Cut_Aera_Dn = (Dn_x, Dn_y)
            Data_ThisPic_AtomNum, Data_ThisPic_AtomRegion, ne = bf.Load_Pic_And_Cal(
                Photo_Ind, Load_Name, File_Mode, Save_Name, SaveFigure_Address + 'All_' + Cam + '_',
                Atom_Center, Cut_Width, len(Cut_Width) - 1, OriImage_Switch,
                Rot_Ang=58, Axis='off', Enclose=Enclose_Switch)
            Data_This_Row[Photo_Ind] = (Data_ThisPic_AtomNum, Data_ThisPic_AtomRegion)
            Atom_Clouds.append(np.array(Data_ThisPic_AtomNum))
            Data_ThisPic_AtomNum, Data_ThisPic_AtomRegion, Up_Region = bf.Load_Pic_And_Cal(
                Photo_Ind, Load_Name, File_Mode, Save_Name, SaveFigure_Address + 'Up_' + Cam + '_',
                Atom_Center, Cut_Width, len(Cut_Width) - 1, OriImage_Switch,
                Rot_Ang=58, Axis='off', Cut_Aera=Cut_Aera_Up, Enclose=Enclose_Switch)
            Up_Regions.append(Up_Region)
            # Data_ThisPic_AtomNum, Data_ThisPic_AtomRegion, Dn_Region = bf.Load_Pic_And_Cal(
            #     Photo_Ind, Load_Name, File_Mode, Save_Name, SaveFigure_Address+'Down_'+Cam+'_',
            #     Atom_Center, Cut_Width, len(Cut_Width)-1, OriImage_Switch,
            #     Rot_Ang=58, Axis='off', Cut_Aera=Cut_Aera_Dn, Enclose=Enclose_Switch)
            Dn_Regions.append(Up_Region)
            figsize = (6.5, 10)
            Top_Up = 0.98
            Bottom_Up = 0.039
            Left_Up = 0.015
            Right_Up = 0.985
            Hspace_Up = 0.2
            Wspace_Up = 0.2
            Top_Dn = 0.98
            Bottom_Dn = 0.039
            Left_Dn = 0.015
            Right_Dn = 0.985
            Hspace_Dn = 0.2
            Wspace_Dn = 0.2

        Photo_Ind += 1
    Up_Average = np.average(np.array(Up_Regions), 0)
    Dn_Average = np.average(np.array(Dn_Regions), 0)
    #    input("Averages:\n")
    theta = np.linspace(-np.pi, np.pi, 1000)
    fig = plt.figure(10000, figsize=figsize)
    plt.imshow(Up_Average, vmin=-0.1, vmax=0.5, cmap='jet')
    plt.subplots_adjust(
        top=Top_Up, bottom=Bottom_Up, left=Left_Up,
        right=Right_Up, hspace=Hspace_Up, wspace=Wspace_Up)
    for Atom_Ind in range(len(Cut_Width) - 1):
        if Enclose_Switch == 0:
            break
        circlex = Cut_Width[Atom_Ind] / 2 * np.cos(theta) + Atom_Center[Atom_Ind][0] - Up_x[0]
        circley = Cut_Width[Atom_Ind] / 2 * np.sin(theta) + Atom_Center[Atom_Ind][1] - Up_y[0]
        plt.plot(circlex, circley, 'k')
    plt.axis('off')
    fig.savefig(SaveFigure_Address + Cam + '_' + Photo + r'Up_Average_Large' + '.png')
    fig.savefig(SaveFigure_Address + Cam + '_' + Photo + r'Up_Average_Large' + '.pdf')
    plt.pause(0.5)

    fig = plt.figure(10001, figsize=figsize)
    plt.imshow(Dn_Average, vmin=-0.1, vmax=0.5, cmap='jet')
    plt.subplots_adjust(
        top=Top_Dn, bottom=Bottom_Dn, left=Left_Dn,
        right=Right_Dn, hspace=Hspace_Dn, wspace=Wspace_Dn)
    plt.colorbar()
    fig.savefig(SaveFigure_Address + Cam + '_' + Photo + r'Up_Average_no_circle_Large' + '.png')
    fig.savefig(SaveFigure_Address + Cam + '_' + Photo + r'Up_Average_no_circle_Large' + '.pdf')
    plt.show()
    
    fig = plt.figure(10002, figsize=figsize)
    im = plt.imshow(Dn_Average[258:559, 188:489], vmin=-0.1, vmax=0.5, cmap='jet')
    plt.colorbar(im, fraction=0.0455, pad=0.04)
    plt.axis('off')
    fig.savefig(SaveFigure_Address + Cam + '_' + Photo + r'Up_Average_no_circle' + '.png')
    fig.savefig(SaveFigure_Address + Cam + '_' + Photo + r'Up_Average_no_circle' + '.pdf')
    plt.show()
    
    #    for Atom_Ind in range(5, 11):
    #        circlex = Cut_Width[Atom_Ind]/2*np.cos(theta)+Atom_Center[Atom_Ind][0]-Up_x[0]
    #        circley = Cut_Width[Atom_Ind]/2*np.sin(theta)+Atom_Center[Atom_Ind][1]-Up_y[0]
    #        plt.plot(circlex, circley)
    #    plt.axis('off')
    #    fig.savefig(SaveFigure_Address+Cam+'_'+Photo+r'Dn_Average'+'.png')
    #    fig.savefig(SaveFigure_Address+Cam+'_'+Photo+r'Dn_Average'+'.pdf')
    #    plt.pause(0.5)
    #    if Cam == '22834':
    #        input("alll")
    Data_Every_Row[Row_Num] = (Data_This_Row, Up_Average, Dn_Average)
    AtomClouds_Every_Row[Row_Num] = np.average(np.array(Atom_Clouds), 0)

time_end = time.time()
print('Time spent: ', time_end - time_start, 's')
